from tac3d_toolkit.geometry import get_normal_field
from tac3d_toolkit.visualize import plot_vector_field, plot_scalar_field_animation, plot_vector_field_animation
from tac3d_toolkit.basic import extract_tacdata_from_frame, unzip_PDF
from tac3d_toolkit.force import get_xy_force, get_force_norm, get_average_normal, force_normal_decompose, get_fr_normal_decomposition
from tac3d_toolkit.slippage import get_C_strains
from tac3d_toolkit.contact import get_contact_region
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# create a 20 * 20 grid of points
P = np.array([[x, y, 1-x**2 - y** 2] for x in np.linspace(-1, 1, 20) for y in np.linspace(-1, 1, 20)])

F = np.array([[0, 0, 1] for x in np.linspace(-1, 1, 20) for y in np.linspace(-1, 1, 20)])
Fn, Ft = force_normal_decompose(P, F)
avg_normal = get_average_normal(P, F)
Fnr, Ftr = get_fr_normal_decomposition(P, F)
Fr = Fnr + Ftr
print(Fnr+Ftr, Fnr, Ftr, avg_normal)
normal = get_normal_field(P)
position = P.reshape(-1,3)
array = F.reshape(-1,3)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 由位置场散点拟合表面形貌并绘制

X = position[:, 0].reshape(20, 20)
Y = position[:, 1].reshape(20, 20)
Z = position[:, 2].reshape(20, 20)
ax.plot_surface(X, Y, Z, alpha=0.3)
# 画向量场
ax.quiver(position[:,0], position[:,1], position[:,2],
            array[:,0], array[:,1], array[:,2],
            length=0.1, normalize=False, color=[1,0,0],label='Force Vectors',alpha=0.5)
# # 画法向量场
# ax.quiver(position[:,0], position[:,1], position[:,2],
#             normal[:,0], normal[:,1], normal[:,2],
#             length=0.1, normalize=False, color=[0,0,0],label='Normal Vectors',alpha=0.1)
# 画微元上法向力
# ax.quiver(position[:,0], position[:,1], position[:,2],
#             Fn[:,0], Fn[:,1], Fn[:,2],
#             length=0.1, normalize=False, color=[0,0,1],label='Normal Force Vectors',alpha=0.6)
# # 画微元上切向力
# ax.quiver(position[:,0], position[:,1], position[:,2],
#             Ft[:,0], Ft[:,1], Ft[:,2],
#             length=0.1, normalize=False, color=[1,0,0],label='Shear Force Vectors',alpha=0.6)

ax.legend()
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_aspect('equal')
# reduce the margin of the plot
plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
plt.title("Average Normal Vector in constant force field")
plt.show()
